#!/usr/bin/env python
"""

    run_svm_c2h2_classify.py

"""
################################################################################
#
#   run_muscle
#
#
#   Copyright (C) 2009 Leo Goodstadt
#
#   This program is free software; you can redistribute it and/or
#   modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation; version 2
#   of the License
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#################################################################################

import sys, os
from collections import defaultdict
import json
#import popen2
import subprocess
import tempfile
import string
import time
import random
from itertools import product
#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888


#_________________________________________________________________________________________
#
#   get_feature_strings
#
#_________________________________________________________________________________________

def get_feature_strings (amino_acid_sequence):
    """
        Get specification strings for svm_classify

            to estimate the affinities between one zinc finger ("amino_acid_sequence")
            and all of the possible corresponding DNA at each of the 4 positions.

            Since there are 4 nucleotide types ACGT, There are
            4 x 4 x 4 x 4  = 256 interaction affinities we need to estimate.

            The 256 possibilities are enumerated in order,
                from AAAA, AAAC .... TTTG, TTTT

            To convert from index to the AA,

                Base 1 = (i / 64) % 4
                Base 2 = (i / 16) % 4
                Base 3 = (i /  4) % 4
                Base 4 = i        % 4

                3 - index for any base for complimentary strand

            For example:
                from itertools import product
                quad_to_seq = ["".join(f) for f in product(*["ACGT"] * 4)
                for i in range(256):
                    print i, quad_to_seq[i], "ACGT"[(i / 64) % 4], "ACGT"[(i / 16) % 4], "ACGT"[(i /  4) % 4], "ACGT"[i        % 4]
                    print i, quad_to_seq[i], "ACGT"[3 - (i / 64) % 4], "ACGT"[3 - (i / 16) % 4], "ACGT"[3 - (i /  4) % 4], "ACGT"[3 - i  % 4]


            From the zinc finger model:

                In each zinc finger, the
                four amino acid positions that contact DNA are called 1, 2,
                3, and 6 (as numbered from the start of the alpha-helix). We use
                the this model to represent each protein-DNA complex by a feature
                vector x = {x_abc}, where x_abc=1 for every amino acid 'a'
                from {A,C,D,...,Y} interacting with base 'b' from {a,c,g,t} at
                contact position 'c'. This representation scheme leads to a
                feature space containing 320 dimensions representing all possible
                abc combinations (20 amino acids x 4 bases x 4 contacts).
                Therefore, each 'abc' combination can be represented as a number
                in (1,320) range according to the conversion table:

                Conversion table from SVM features to amino acid - nucleotide contacts.
                Canonials contacts are listed as (see Fig 1 at Persikov et al., 2008):
                01 - between amino acid a6 and nucleotide b1
                02 - between amino acid a3 and nucleotide b2
                03 - between amino acid a1 and nucleotide b3
                04 - between amino acid a2 and nucleotide b4'

                  1 01Aa
                  2 01Ac
                  3 01Ag
                  4 01At
                  5 01Ca
                ...
                316 04Wt
                317 04Ya
                318 04Yc
                319 04Yg
                320 04Yt

                Finally, for the SVM-light input, each tested zinc finger-DNA pair should
                be represented as a feature vector recorded as a single line in the
                example_file. It is of the following format:

                <line> .=. <target> <feature>:<value> <feature>:<value> ... <feature>:<value> # <info>
    """
    if len(amino_acid_sequence) != 4:
        raise Exception("Expected a four amino acid string [%s]" % (amino_acid_sequence,) )

    valid_aa = dict(zip("ACDEFGHIKLMNPQRSTVWY", range(20)))


    for aa in amino_acid_sequence:
        if not aa in valid_aa:
            raise Exception("Non-conforming amino acid [%s]" % (aa,))


    svm_values = []


    # change order to a6, a3, a0 and a2
    amino_acid_sequence = [amino_acid_sequence[aa] for aa in [3, 2, 0, 1]]

    #
    #   For each amino acid a6, a3, a0, permute ACGT
    #
    for aa_pos, aa in enumerate(amino_acid_sequence[:-1]):
        aa_index = valid_aa[aa]
        svm_values.append([((aa_pos) * 20 * 4 + 4 *(aa_index) + (acgt) + 1) for acgt in range(4)])

    #
    #   For amino acid a2, permute TGCA because a2 binds to b4' on the complementary strand
    #       This is the ONLY place we want to deal with complementary bases
    #
    aa_pos, aa = 3, amino_acid_sequence[-1]
    aa_index = valid_aa[aa]
    svm_values.append([((aa_pos) * 20 * 4 + 4 *(aa_index) + (acgt) + 1) for acgt in reversed(range(4))])

    #print svm_values

    feature_strings = []
    for features in product(*svm_values):
        #0 3:1 59:1 94:1 131:1 # test
        feature_strings.append("0 %d:1 %d:1 %d:1 %d:1\n" % features)

    return feature_strings


#_________________________________________________________________________________________
#
#   get_temp_file_name
#
#_________________________________________________________________________________________
def get_temp_file_name (**args):
    """
    Creates temporary file and returns file name
    Set optionally
        mode='w+b'
        bufsize=-1
        suffix=''
        prefix='tmp'
        dir=None
    """
    temp_file = tempfile.NamedTemporaryFile(delete=False, **args)
    temp_file_name = str(temp_file.name)
    temp_file.close()
    return temp_file_name


#_________________________________________________________________________________________
#
#   _cleanup_files
#
#_________________________________________________________________________________________
from contextlib import contextmanager
@contextmanager
def _cleanup_files (file_names):
    """
    Cleanup temp files
    """
    yield
    for f in file_names:
        try:
            #print f
            #sys.stdout.write(open(f).read())
            #print "\n" *2

            os.unlink(f)

            pass
        except:
            pass

#_________________________________________________________________________________________
#
#   max_svm_c2h2_scores_per_16
#
#_________________________________________________________________________________________
from collections import namedtuple
t_score_sequence = namedtuple("t_score_sequence", "score quad")
def max_svm_c2h2_scores_per_16 (values):
    """
        Return affinity scores from svm as a list of scores / 4-base nucleotides

        [index base 1][index of base4] = max(scores for any base2 / base3), B1-4

        base index  : ACGT = 0123

    """
    #quad_to_seq = ["".join(f) for f in product(*["ACGT"] * 4)]
    values_seq_by_b1_next_b1 = [[[] for i in range(4)] for j in range(4)]
    for i in range(256):
        b1      = (i / 64)
        next_b1 = i % 4
        values_seq_by_b1_next_b1[b1][next_b1].append(t_score_sequence(values[i], i))
    for ii in range(4):
        for jj in range(4):
            values_seq_by_b1_next_b1[ii][jj] = max(values_seq_by_b1_next_b1[ii][jj])

    return values_seq_by_b1_next_b1



#_________________________________________________________________________________________
#
#   run_svm_c2h2_classify
#
#_________________________________________________________________________________________
def run_svm_c2h2_classify (amino_acids, svm_classify_bin = None, svm_model_file_name = None, shuffle_randomly = False):
    """
    Calculates the C2H2 binding of zinc fingers to dna sequence
        http://svmlight.joachims.org/
        See http://bioinformatics.oxfordjournals.org/cgi/content/short/25/1/22
        Anton V. Persikov, Robert Osada, and Mona Singh
        Predicting DNA recognition by Cys2His2 zinc finger proteins
        Bioinformatics 2009 25: 22-29.

    Returns a list of numbers representing matches
            0 -255
            enumerating from AAAA, AAAC .... TTTG, TTTT
        where AAAA etc are the residues for b1->b4 all on the same strand

    N.B.
        This is NOT b'4 but b4



    """
    if not svm_classify_bin:
        svm_classify_bin = "/net/cpp-group/tools/svm_classify_C2H2_binding/svm_classify"
    if not svm_model_file_name:
        svm_model_file_name = "/net/cpp-group/tools/svm_classify_C2H2_binding/SVMp.mod"

    feature_strings = get_feature_strings(amino_acids)
    if shuffle_randomly:
        feature_strings = list(enumerate(feature_strings))
        random.shuffle(feature_strings)
        random_order = [i for (i, s) in feature_strings]
        feature_strings = [s for (i, s) in feature_strings]

    input_file_name = get_temp_file_name()
    input_file = open(input_file_name, "w")
    input_file.write("".join(feature_strings))
    input_file.close()

    output_file_name = get_temp_file_name()

    try:
        # ./svm_classify  example.txt SVMl.mod output.txt
        cmds = " ".join((svm_classify_bin, input_file_name, svm_model_file_name, output_file_name))
        with _cleanup_files([input_file_name, output_file_name]):
            process = subprocess.Popen(cmds,
                                       stderr=subprocess.PIPE,
                                       stdout=subprocess.PIPE,
                                       shell=True
                                       )
            (stdout_str, stderr_str) = process.communicate()
            if process.returncode == 0:
                values = [float (v)  for v in open(output_file_name).read().split("\n") if len(v)]
                if len(values) == 256:
                    if shuffle_randomly:
                        values = sorted(zip(random_order, values))
                        values = [v for (i,v) in values]
                    return values
                err_str = ("svm classify returned the wrong number of values (%d) "
                            "for [%s]\n[%s %s %s]\n%s\n%s\n" % (len(values),
                                                                   amino_acids,
                                                                   cmds,
                                                                   input_file_name,
                                                                   output_file_name,
                                                                   stdout_str,
                                                                   stderr_str))

            #
            #   failed: cleanup and throw exception
            #
            else:

                err_str = ("svm classify failed, returning "
                            "%s for [%s]\n[%s %s %s]\n%s\n%s\n" % (process.returncode,
                                                                   amino_acids,
                                                                   cmds,
                                                                   input_file_name,
                                                                   output_file_name,
                                                                   stdout_str,
                                                                   stderr_str))
        raise Exception(err_str)
    except:
        raise


#_________________________________________________________________________________________
#
#   run_svm_c2h2_classify_piecemeal
#
#_________________________________________________________________________________________
def run_svm_c2h2_classify_piecemeal (amino_acids, svm_classify_bin = None, svm_model_file_name = None):
    """
    Sanity check function:
        same as run_svm_c2h2_classify()
        but runs finger by finger
    """
    if not svm_classify_bin:
        svm_classify_bin = "/net/cpp-group/tools/svm_classify_C2H2_binding/svm_classify"
    if not svm_model_file_name:
        svm_model_file_name = "/net/cpp-group/tools/svm_classify_C2H2_binding/SVMp.mod"

    feature_strings = get_feature_strings(amino_acids)

    input_file_name = get_temp_file_name()
    output_file_name = get_temp_file_name()

    values = []
    for fs in feature_strings:
        with _cleanup_files([input_file_name, output_file_name]):
            input_file = open(input_file_name, "w")
            input_file.write(fs)
            input_file.flush()
            input_file.close()
            try:
                cmds = " ".join((svm_classify_bin, input_file_name, svm_model_file_name, output_file_name))
                process = subprocess.Popen(cmds,
                                           stderr=subprocess.PIPE,
                                           stdout=subprocess.PIPE,
                                           shell=True
                                           )
                (stdout_str, stderr_str) = process.communicate()
                if process.returncode == 0:
                    output_file = open(output_file_name)
                    value = [float (v)  for v in output_file.read().split("\n") if len(v)]
                    output_file.close()
                    if len(value) == 1:
                        values.append(value[0])
                        continue
                    err_str = ("svm classify returned the wrong number of values (%d) "
                                "for [%s]\n[%s %s %s]\n%s\n%s\n" % (len(value),
                                                                       amino_acids,
                                                                       cmds,
                                                                       input_file_name,
                                                                       output_file_name,
                                                                       stdout_str,
                                                                       stderr_str))
                #
                #   failed: cleanup and throw exception
                #
                else:

                    err_str = ("svm classify failed, returning "
                                "%s for [%s]\n[%s %s %s]\n%s\n%s\n" % (process.returncode,
                                                                       amino_acids,
                                                                       cmds,
                                                                       input_file_name,
                                                                       output_file_name,
                                                                       stdout_str,
                                                                       stderr_str))
                raise Exception(err_str)
            except:
                raise
    return values



#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Testing


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
import unittest
class Test_run_svm_c2h2_classify(unittest.TestCase):

    #       self.assertEqual(self.seq, range(10))
    #       self.assert_(element in self.seq)
    #       self.assertRaises(ValueError, random.sample, self.seq, 20)



    def test_function(self):
        """
            test
        """
        import re

        # human prdm9
        dna_sites = [        "AGTGTTGATACA",        # SVDT
                             "AGCTGGCACATT",        # SWHI
                             "AGCTGGGTCACT",        # SWVT
                             "AGCCGGGTCACT",        # SRVT
                             "AGCCGGGTCACT",        # SRVT
                             "AGCTGGGTCACT",        # SWVT
                             "AGCTGGGTCACT",        # SWVT
                             "AGCAATCACAGA",        # SNHR
                             "CGCGATCACAGA",        # RDHR
                             "AGAGATAACAGT",        # RDNS
                             "AGCAATCACAGA",        # SNHR
                             "CGCAATCACAGA",        # RNHR
                             "AGCGATAGCTAT",        # SDSY
                             ]
        from translate_by_codon_old import fast_translate, translate

        #print fast_translate("".join(dna_sites))
        #sys.exit()

        prot_sites = [fast_translate(cdna) for cdna in dna_sites]
        prot_sites = ["DSNS"]

        quad_to_seq = ["".join(f) for f in product(*["ACGT"] * 4)]
        seq_to_quad = dict(zip(quad_to_seq, range(256)))
        from math import log, exp
        log_0109 = log(0.1/0.9)
        beta = 0.24
        def score_to_binding_prob (max_score, score):
            D = exp (log_0109 - (max_score - score) / beta)
            return D / (1 + D)

        def binding_prob_to_entropy (probs):
            sum_probs = sum(probs)
            proportional_probs = [ p / sum_probs for p in probs]
            total_entropy = 2 + sum((p * log(p, 2)) for p in proportional_probs)
            entropies = [(total_entropy * p) for p in proportional_probs]
            return entropies, total_entropy


        def svm_scores_to_entropy (scores):
            max_score = scores[-1]
            probs = [score_to_binding_prob(max_score, s) for s in scores[:-1]]

            score_diffs = [beta * (log(0.1 / 0.9) - log(p/(1-p))) for p  in probs]



            print "\t".join("%.5f" % (max_score - sd)  for sd in score_diffs)

            return binding_prob_to_entropy (probs)



        for prot in reversed(prot_sites):
            values = run_svm_c2h2_classify (prot)
            consensus_quad = seq_to_quad["TACG"]
            #other_quads = [seq_to_quad[s] for s in ["TAAG","TACG","TAGG","TATG"]]
            other_quads = [seq_to_quad[s] for s in ["TACC","TTCG","TCCG","TGCG"]]
            for ii, vv in enumerate(values):
                print quad_to_seq[ii], ii, vv
            print quad_to_seq[consensus_quad], consensus_quad, values[consensus_quad]
            for q in other_quads:
                print quad_to_seq[q], q, values[q]
            scores = [values[q] for q in other_quads]
            scores.append(values[consensus_quad])

            entropies, max_entropy = svm_scores_to_entropy(scores)
            print "\t".join("%.5f" % e for e in entropies)

            #
            #   Sanity check that running svm_classify on multiple fingers at a time
            #       changes nothing
            #
            #values2 = run_svm_c2h2_classify ("".join(valid_aa[0:4]), shuffle_randomly = True)
            #self.assertEqual(values, values2)

            print >>sys.stderr, "%s" % "".join(prot)

            sorted_values = max_svm_c2h2_scores_per_16 (values)
            for ii in range(4):
                for jj in range(4):
                    print "    ACGT"[ii], "ACGT"[jj], sorted_values[ii][jj][1], "%6.3f" % sorted_values[ii][jj][0]


#
#   debug code not run if called as a module
#
if __name__ == '__main__':
    if sys.argv.count("--debug"):
        sys.argv.remove("--debug")
    unittest.main()



